import pyautogui
from ultralytics import YOLO
import win32gui
import win32ui
import win32con
import cv2
import numpy as np
import os
import time
import subprocess
import sys
from yolo2.find_and_perform_action_module import find_and_perform_action
from yolo2.find_class_module import find_class_and_perform_action
from yolo2.moveFb模块 import *
from yolo2.yolo8瞄准模板 import AimbotModule
from yolo2.巡逻模块 import patrol_task
import yolo8瞄准模板
import 检测框模块
import torch

# 用于标记模块是否被加载
MODULE_LOADED = True

# 最新版本的模型路径
MODEL_PATH = r"H:\yol\yolo\v11\runs\detect\train17\weights\best.pt"

# 全局变量，用于存储已加载的模型和窗口句柄
loaded_model = None
game_hwnd = None


def load_model():
    """加载YOLOv8模型（强制使用 CPU）"""
    global loaded_model
    if loaded_model is None:
        pyautogui.FAILSAFE = False
        loaded_model = YOLO(MODEL_PATH)
        # 强制将模型移动到 CPU（忽略 CUDA 可用性）
        loaded_model.to('cpu')
    return loaded_model


def find_game_window():
    """查找标题中包含'梦幻西游'的窗口，避免重复查找"""
    global game_hwnd
    if game_hwnd is None:
        def callback(hwnd, extra):
            window_title = win32gui.GetWindowText(hwnd)
            if '梦幻西游' in window_title:
                extra.append(hwnd)
            return True

        hwnds = []
        win32gui.EnumWindows(callback, hwnds)
        if hwnds:
            game_hwnd = hwnds[0]
    return game_hwnd


def get_window_screenshot(hwnd):
    left, top, right, bottom = win32gui.GetWindowRect(hwnd)
    width = right - left
    height = bottom - top
    hwndDC = win32gui.GetWindowDC(hwnd)
    mfcDC = win32ui.CreateDCFromHandle(hwndDC)
    saveDC = mfcDC.CreateCompatibleDC()
    saveBitMap = win32ui.CreateBitmap()
    saveBitMap.CreateCompatibleBitmap(mfcDC, width, height)
    saveDC.SelectObject(saveBitMap)
    saveDC.BitBlt((0, 0), (width, height), mfcDC, (0, 0), win32con.SRCCOPY)
    signedIntsArray = saveBitMap.GetBitmapBits(True)
    img = np.frombuffer(signedIntsArray, dtype='uint8')
    img.shape = (height, width, 4)
    win32gui.DeleteObject(saveBitMap.GetHandle())
    saveDC.DeleteDC()
    mfcDC.DeleteDC()
    win32gui.ReleaseDC(hwnd, hwndDC)
    return cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)


def draw_detection_boxes(screenshot, results, model):
    for result in results:
        boxes = result.boxes.numpy()  # 直接获取 CPU 上的 numpy 数组（无需额外转换）
        for box in boxes:
            class_id = int(box.cls[0])
            conf = round(box.conf[0], 2)
            class_name = model.names[class_id]
            x1, y1, x2, y2 = box.xyxy[0].astype(int)
            cv2.rectangle(screenshot, (x1, y1), (x2, y2), (0, 255, 0), 2)
            label = f"({class_name} + ID: {class_id})"
            cv2.putText(screenshot, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
    return screenshot


def save_screenshot(screenshot, save_dir):
    """保存截图到指定目录"""
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    save_path = os.path.join(save_dir, f"screenshot_{int(time.time())}.png")
    cv2.imwrite(save_path, screenshot)
    print(f"游戏窗口截图已保存到 {save_path}")


def run_subprocess(script_path):
    """运行子进程脚本"""
    subprocess.run([sys.executable, script_path], check=False)


def start_detection():
    hwnd = find_game_window()
    if hwnd is None:
        print("未找到标题中包含'梦幻西游'的窗口")
        return

    cv2.namedWindow("Detection Window", cv2.WINDOW_NORMAL)
    model = load_model()

    while True:
        screenshot = get_window_screenshot(hwnd)
        results = model(screenshot)  # 直接在 CPU 上推理
        screenshot_with_boxes = draw_detection_boxes(screenshot, results, model)

        cv2.imshow("Detection Window", screenshot_with_boxes)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cv2.destroyAllWindows()


if __name__ == "__main__":
    start_detection()